import asyncio
from xhs_auto_poster.__001__langgraph_more_nodes.agent_state import AgentState
from langchain_core.messages import HumanMessage
from xhs_auto_poster.__002__fastapi.msg_queue import MSG_TYPE, msg_queue
from xhs_auto_poster.common.llm import my_llm  # 假设 my_llm 是 ChatOpenAI / ChatGLM 等支持 astream 的实例


async def llm_direct_out_node(state: AgentState):
    print("开始生成直接用户回答")
    await msg_queue.put({"type": MSG_TYPE.THINK.name, "content": "开始生成直接用户回答"})
    await asyncio.sleep(0.1)

    # 获取用户输入
    user_input = state["input"]

    # 构建提示词
    prompt = f"你是一个旅行专家，请根据用户输入，生成回答。\n用户输入：{user_input}"

    # 累积最终结果
    final_answer = []

    # ---- 关键：使用 astream 进行流式输出 ----
    async for chunk in my_llm.astream([HumanMessage(content=prompt)]):
        # chunk 可能是 ChatGenerationChunk / AIMessageChunk
        delta = getattr(chunk, "content", None)
        if delta:
            final_answer.append(delta)
            # 每次追加新的文字就推送给前端
            await msg_queue.put({
                "type": MSG_TYPE.ANSWER.name,
                "content": delta  # 也可以拼接 ''.join(final_answer) 看前端需求
            })

    # 拼接完整答案
    model_answer = "".join(final_answer).strip()

    # 存入 state
    state["direct_out"] = model_answer
    state['output'] = model_answer
    print("完成生成直接用户回答")
    await msg_queue.put({"type": MSG_TYPE.THINK.name, "content": "完成生成直接用户回答"})
    await asyncio.sleep(0.1)

    return state
